open Ir

module OperandKey = struct
  type t = operand

  let equal a b =
    match (a, b) with
    | Var x, Var y -> String.equal x y
    | Reg x, Reg y -> String.equal x y
    | Imm x, Imm y -> x = y
    | _, _ -> false

  let hash = function
    | Var x -> Hashtbl.hash ("var", x)
    | Reg x -> Hashtbl.hash ("reg", x)
    | Imm i -> Hashtbl.hash ("imm", i)
end

module OperandHashtbl = Hashtbl.Make (OperandKey)

(* Reg 或者 Var 的活跃区间 *)
(* Var 在逻辑上和 Reg 等价 *)
type interval = {
  name : operand;
  mutable start : int;
  mutable end_ : int;
  is_param : bool;
}

let spill_map : (operand, int) Hashtbl.t = Hashtbl.create 32
let stack_slot_counter = ref 0

(* 假设有 K 个可用物理寄存器 *)
let k_registers = 8 + 4 + 11

(* t4, t5, t6 被保留用来处理 spill 变量 *)
(* s0 用于充当帧指针 *)
let physical_registers =
  [|
    "a0";
    "a1";
    "a2";
    "a3";
    "a4";
    "a5";
    "a6";
    "a7";
    "t0";
    "t1";
    "t2";
    "t3";
    (* "t4"; *)
    (* "t5"; *)
    (* "t6"; *)
    (* "s0"; *)
    "s1";
    "s2";
    "s3";
    "s4";
    "s5";
    "s6";
    "s7";
    "s8";
    "s9";
    "s10";
    "s11";
  |]

let print_reg_alloc_result reg_map : unit =
  Printf.printf "=== Register Allocation Result ===\n";
  OperandHashtbl.iter
    (fun op reg ->
      let name =
        match op with
        | Reg r -> Printf.sprintf "Reg %s" r
        | Var v -> Printf.sprintf "Var %s" v
        | Imm i -> Printf.sprintf "Imm %d" i
      in
      Printf.printf "%-10s -> %s\n" name reg)
    reg_map

(* 1. 构建活跃区间 *)
let build_intervals (f : ir_func_o) : interval list =
  let reg_intervals = Hashtbl.create 128 in
  let block_pos = Hashtbl.create 64 in
  let counter = ref 0 in

  List.iter
    (fun block ->
      incr counter;
      Hashtbl.add block_pos block.label !counter)
    f.blocks;

  let is_param_operand op =
    match op with Var name -> List.mem name f.args | _ -> false
  in

  let add_live v pos =
    let op = v in
    match Hashtbl.find_opt reg_intervals op with
    | Some i ->
        i.start <- min i.start pos;
        i.end_ <- max i.end_ pos
    | None ->
        let i =
          { name = op; start = pos; end_ = pos; is_param = is_param_operand op }
        in
        Hashtbl.add reg_intervals op i
  in

  List.iter
    (fun block ->
      let pos = Hashtbl.find block_pos block.label in
      OperandSet.iter (fun v -> add_live v pos) block.live_in;
      OperandSet.iter (fun v -> add_live v pos) block.live_out)
    f.blocks;

  Hashtbl.fold (fun _ itv acc -> itv :: acc) reg_intervals []

(* 2. 分配物理寄存器 *)
let linear_scan_allocate (intervals : interval list) (print_allocation : bool) =
  let intervals = List.sort (fun a b -> compare a.start b.start) intervals in
  let active : interval list ref = ref [] in

  let reg_map = OperandHashtbl.create 32 in
  (* 当前活跃变量分配 *)
  let alloc_map = OperandHashtbl.create 512 in
  (* 所有分配历史，最终用于输出 *)

  let expire_old_intervals current =
    (* 重新设置 active 列表 *)
    active :=
      List.filter
        (fun itv ->
          if itv.end_ >= current.start then true
          else (
            (* 释放该变量使用的寄存器 *)
            OperandHashtbl.remove reg_map itv.name;
            false))
        !active
  in

  List.iter
    (fun itv ->
      (* itv 还没有在 active 里面 *)
      expire_old_intervals itv;

      if itv.is_param then (
        OperandHashtbl.add reg_map itv.name "__SPILL__";
        OperandHashtbl.replace alloc_map itv.name "__SPILL__")
      else if List.length !active = k_registers then (
        (if
           (* spill: 当前无法分配物理寄存器，简单略过或标记 *)
           (* TODO: spill 到栈上和对 active 区间的处理 *)
           print_allocation
         then
           let name =
             match itv.name with Reg r | Var r -> r | Imm _ -> "imm"
           in
           Printf.printf "Spill: %s\n" name);
        OperandHashtbl.add reg_map itv.name "__SPILL__";
        (* 明确标记为 spill *)
        OperandHashtbl.replace alloc_map itv.name "__SPILL__")
      else
        (* 分配一个可用寄存器 *)
        (* Hashtbl -> list *)
        let used_regs =
          OperandHashtbl.fold (fun _ r acc -> r :: acc) reg_map []
        in
        let avail =
          List.find
            (fun r -> not (List.mem r used_regs))
            (Array.to_list physical_registers)
        in
        OperandHashtbl.add reg_map itv.name avail;
        OperandHashtbl.replace alloc_map itv.name avail;
        (* let name =
          match itv.name with
          | Var v -> "Var " ^ v
          | Reg v -> "Reg " ^ v
          | _ -> "imm"
        in
        Printf.printf "Allocate virtual %s with physical reg %s\n" name avail; *)
        active := itv :: !active)
    intervals;
  if print_allocation then print_reg_alloc_result alloc_map;
  alloc_map
